#pragma once

#include <tr1/unordered_map>
//#include <exception>
#include <algorithm>
#include <stack>
#include <vector>
#include <map>
#include <iostream>
#include "bitstream.h"

using namespace std;

typedef map<ui64, float> TFreqsMap;


// <symb, freq>
template <typename T> class TFreqTable;
template <typename T> class THuffmanCodec;

// 1, 2 and 4 values allowed
const ui32 MAX_BYTES_IN_SYMBOL = 4;

/*
 * THuffmanCodec
 */
template <typename T>
class THuffmanCodec
{
public:	
    enum {
        // invalide code - indicates some code that will be passed over
		INVALID_CODE = ui64(-1)
        // ecape symbol for escape code, which will be put at the end of Decoded file
		,ESCAPE_SYMBOL = ui64(-2)
		// number of symbols above values of T
		,NUM_OF_SPEC_SYMBOLS = 2
		// bytes in symbol
		,LEN_OF_SYMBOL = sizeof(T)
		// total number of symbils
		,NUM_OF_SYMBOLS = (1<<(8*LEN_OF_SYMBOL)) + NUM_OF_SPEC_SYMBOLS
		// max number of nodes
		,NUM_OF_NODS = (NUM_OF_SYMBOLS*2-1)
		// max len of code in bytes
		,MAX_CODE_LEN_IN_BYTES = (NUM_OF_SYMBOLS/8)
		// amount of bits reading at once during decoding
        ,DECODE_JUMP_LEN = 11
        // max num of nodes saved in start jump table
        ,NODES_STORED = 2 // only 1 and 2 implemented
    };

private:
    typedef ui32 TNodeID;
    struct TEntry
    {
        union {
            ui64 Code;
            struct {
                TNodeID Left;
                TNodeID Right;
            } Composite;
        } Data;
        TEntry(ui64 code)
        {
            Data.Code = code;
        }
        TEntry(TNodeID left, TNodeID right)
        {
            Data.Composite.Left = left;
            Data.Composite.Right = right;
        }
    };
    // comparator
    struct TCmpEntries
    {
		const TFreqsMap& Freqs;
        TCmpEntries(const TFreqsMap& freqs)
            : Freqs(freqs)
        {}
        bool operator ()(const TEntry& a, const TEntry& b) const
        {
            return Freqs.find(a.Data.Code)->second < Freqs.find(b.Data.Code)->second;
        }
    };
    typedef vector<TEntry> THuffmanTree;
    // <symb, <code, len>>
    typedef tr1::unordered_map<ui64,std::pair<ui64, ui32> > TIndexMap;
    struct TStartJumpResult {
        enum {SIZE=NODES_STORED};
        ui32 Len;
        TNodeID Node;
        TNodeID Node2;
    };
    typedef vector<TStartJumpResult> TStartJumpTable;

    // node -> symb if leaf
    // node -> node_left, node_right    
    THuffmanTree CodesTree;
    TNodeID TreeHead;
    // symb -> (code, length)
    TIndexMap DirectIndex;
    TStartJumpTable StartJumpTable;

public:
    THuffmanCodec () : TreeHead(0) {}
    THuffmanCodec (const TFreqsMap& freqs) : TreeHead(0) {
        Initialize (freqs);
    }
    ~THuffmanCodec() {
        //if (IsInitialized()) {}
    }
    bool IsInitialized() {
        return (TreeHead);
    }

    template < typename TSource, typename TDataWriter >
    void Code(TTDataReader<TSource>& byteInput, TTBitWriter<1, TDataWriter>& bitOutput) {

		const ui32 size = byteInput.Size()/sizeof(T);
		union {
			ui64 as64;
			T asT;
		} symb;
		symb.as64 = 0;

        try {
            for (ui32 i = 0; i < size; i++) {
				byteInput.ReadInto(symb.asT);
				bitOutput.Write(DirectIndex[symb.as64].first, DirectIndex[symb.as64].second);
            }
            bitOutput.Write(DirectIndex[ESCAPE_SYMBOL].first, DirectIndex[ESCAPE_SYMBOL].second);
        } catch (TException& e) {
            cout << e.what() << endl;
        }
        bitOutput.Flush();
    }

    template<typename TStorage, typename TDataReader>
    void Decode(TTBitReader<TDataReader>& bitInput, TTDataWriter<1, TStorage>& byteOutput) {
        //set code to nil
        ui64 code = 0;
        //set code remains len to nil
        ui32 remainsLen = 0;
        // we seek for symbzz
		union {
            ui64 as64;
            T asT;
		} symb;
        TIMER_CODE (TStopwatch watch;)

		try {
			while(true) {
                //read bits to complete remained bits to a full jump len
                if (remainsLen) {
					const ui32 bitsToRead = DECODE_JUMP_LEN - remainsLen;
                    ui64 newBits = 0;
                    bitInput.Read64(newBits, bitsToRead);
                    //form code
                    code >>= bitsToRead;
                    code |= newBits<<remainsLen;
                } else {
                    bitInput.Read64(code, DECODE_JUMP_LEN);
                }
                //get start jump result
                TIMER_CODE (watch.ResetTsc();)
                const TStartJumpResult jumpResult = StartJumpTable[code];
                //#define jumpResult StartJumpTable[code]
                TIMER_CODE (watch.Count();)
                //set remains len to result len
                remainsLen = jumpResult.Len;
                if (IsNodeLeaf(jumpResult.Node)) {
					symb.as64 = CodesTree[jumpResult.Node].Data.Code;
                } else {
                    //set current node to result node
                    TNodeID currentNode = jumpResult.Node;
                    //till found leaf
                    do {
                        //
                        ui64 dirBit;
                        //reading bit
                        bitInput.ReadABit(dirBit);
                        //move through a tree
                        currentNode = dirBit ? CodesTree[currentNode].Data.Composite.Right
										   : CodesTree[currentNode].Data.Composite.Left;
                    } while (IsNodeLeaf(currentNode) == false);
					symb.as64 = CodesTree[currentNode].Data.Code;
                }

				if (symb.as64 == ESCAPE_SYMBOL) {
                    break;
                } else {
                    byteOutput.WriteWhole(symb.asT);
                }
                // for node2
                if (IsNodeLeaf(jumpResult.Node2)) {
                    symb.as64 = CodesTree[jumpResult.Node2].Data.Code;
                    if (symb.as64 == ESCAPE_SYMBOL) {
                        break;
                    } else {
                        byteOutput.WriteWhole(symb.asT);
                    }
                }
            }
		} catch (TException& e) {
			cout << "Decoding fails. " <<  e.what() << endl;
		} catch (string &e) {
			cout<<"Decoding fails: " <<  e << endl;
		}
        TIMER_CODE (cout<<"Avg: " <<watch.GetAverageTsc()<<endl;)
    }

	template<typename TStorage, typename TSource>
	void bybitDecode(TTBitReader<TSource>& bitInput, TTBitWriter<1, TStorage>& bitOutput) {

        const ui32 bitsOfT = sizeof(T)*8;
        ui64 symb;
        ui32 currentNode;
        ui64 part;

        TIMER_CODE (TStopwatch watch;)
        try {
            while (true) {
				currentNode = TreeHead;
                // seek for leaf
                while (!IsNodeLeaf(currentNode)) {
					//bitInput.Read(part,1);
					bitInput.ReadABit(part);
					currentNode = part ? CodesTree[currentNode].Data.Composite.Right
									   : CodesTree[currentNode].Data.Composite.Left;
                }
                // we found leaf
                symb = CodesTree[currentNode].Data.Code;
                TIMER_CODE (watch.ResetTsc();)
                if (symb == ESCAPE_SYMBOL) {
                    break;
                } else {
                    bitOutput.Write(symb,bitsOfT);
				}
	TIMER_CODE (watch.Count();)
			}
        } catch (TException& e) {
            cout << "Decoding fails. " <<  e.what() << endl;
        }
	TIMER_CODE (cout<<"Avg: " <<watch.GetAverageTsc()<<endl;)
        bitOutput.Flush();
    }

	void oldDecode(const char* start, const char* end, vector<T>* decoded)
	{
		TBitInput<1> in(&start, &end);
		TBitInput<0> exc(&end, &start);
		TIMER_CODE (TStopwatch watch;)
		while (start < end) {
            TNodeID current = TreeHead;
			while (current >= TreeHead) {
				ui64 part = 0;
				if (!in.Read(part, 1)) {
					return;
				}
				current = part ? CodesTree[current].Data.Composite.Right
							   : CodesTree[current].Data.Composite.Left;
			}
			ui64 code = CodesTree[current].Data.Code;
			TIMER_CODE (watch.ResetTsc();)
			if (code == INVALID_CODE) {
				break;
			} else {
				union {
					ui64 val;
					T val2;
				} gccSuck;
				gccSuck.val = code;
				decoded->push_back(gccSuck.val2);
			}
            TIMER_CODE (watch.Count();)
        }
        TIMER_CODE (cout<<"Avg: " <<watch.GetAverageTsc()<<endl;)
    }

private:
    void Initialize (const TFreqsMap& freqs) {
        DEBUG_ASSERT(freqs.size() <= NUM_OF_SYMBOLS);
        if (freqs.empty())
            return;

        // num of significent symbols (leafs)
        const ui32 leavesCount = freqs.size();


        // put codes of significent symbols into CodesTree and sort it
        TCmpEntries cmp(freqs);
        CodesTree.clear();
        //CodesTree.reserve(NUM_OF_NODS);
		typename TFreqsMap::const_iterator iter = freqs.begin();
        while (!(iter == freqs.end())) {
            CodesTree.push_back(iter->first);
            iter++;
        }
        std::sort(CodesTree.begin(), CodesTree.end(), cmp);

        /*
         * Building of a huffman's tree.
         */
        // map of shares for composite nodes
        TFreqsMap compositeShares;
        TCmpEntries compositeCmp(compositeShares);
        // unprocessed leaves list's head elem index
        TNodeID lheadIndex = 0;
        // first level composite nodes list's head elem index
        TNodeID cheadIndex = leavesCount;
        for (ui32 j = 0; j<leavesCount-1; j++) {
            TNodeID join[2];
            float share = 0;
            // choose two lesser nodes from list of leaves(CodesTree starting from lheadIndex) or
            // from list of composite nodes (CodesTree starting from cheadIndex)
            for (int i = 0; i < 2; i++) {
                // frequency(share) of current entry(leaf)
                const float lshare = lheadIndex < leavesCount ? freqs.find(CodesTree[lheadIndex].Data.Code)->second : -1;
                // share of composite node of a tree,  at first -1: cheadIndex = leavesCount = CodesTree.size()
                // afterwards we will add some composite nodes to CodesTree
                const float cshare = cheadIndex < CodesTree.size() ? compositeShares[CodesTree[cheadIndex].Data.Code] : -1;
                // choose lesser one
                const bool leaf = lheadIndex < leavesCount && (cheadIndex >= CodesTree.size() || lshare < cshare);
                join[i] = leaf ? lheadIndex++ : cheadIndex++;
                share += leaf ? lshare : cshare;
            }
            // add new composite node
            TEntry comp(join[0], join[1]);
            CodesTree.push_back(comp);
            compositeShares[CodesTree[CodesTree.size() - 1].Data.Code] = share;
            // sort list of first level composite nodes ?? it will change indexes, isn't it?
            if (cheadIndex < CodesTree.size()) {
                std::sort((typename vector<TEntry>::iterator)&CodesTree[cheadIndex], CodesTree.end(), compositeCmp );
            }
        }
        // reverse order of composite nodes, put the heviest first
        std::reverse((typename vector<TEntry>::iterator)&CodesTree[leavesCount], CodesTree.end());
        for (ui32 i = leavesCount; i < CodesTree.size(); i++) {
            if (CodesTree[i].Data.Composite.Left >= leavesCount)
                CodesTree[i].Data.Composite.Left = CodesTree.size() - (CodesTree[i].Data.Composite.Left - leavesCount + 1);
            if (CodesTree[i].Data.Composite.Right >= leavesCount)
                CodesTree[i].Data.Composite.Right = CodesTree.size() - (CodesTree[i].Data.Composite.Right - leavesCount + 1);
        }
        //DEBUG_ASSERT(CodesTree.size() == NUM_OF_NODS);

        DirectIndex.clear();
        TreeHead = leavesCount;
        std::stack<TNodeID> parents;
        ui32 code = 0;
        ui8 bitsCount = 0;
        ui32 current = TreeHead;
        // without it algorithm never go to the right part of a tree
        parents.push(TreeHead);

		DEBUG_CODE(PrintTree();)
        /*
         * Setting huffman's code values for leaves.
         */
        while (true) {
            MASSERT(bitsCount < 64, "Could not build a tree.");
            // if leaf, save code, pop position from the stack and go right
            if (IsNodeLeaf(current)) {
                DirectIndex[CodesTree[current].Data.Code] = TIndexMap::mapped_type(code, bitsCount);
				DEBUG_CODE( PrintNode(current);  )
                // go down one step and for the first 'left turn' and change it to the 'right'
                bool lastBit;
                do {
                    current = parents.top();
                    parents.pop();
                    // here I have changed code order
                    //bitsCount--;
                    //lastBit = code & 1;
                    //code >>= 1;
                    lastBit = code & (1<<--bitsCount);
                    code &=  ~(1<<bitsCount);
                } while (lastBit);
                if (parents.empty())
                    break;
                parents.push(current);
                current = CodesTree[current].Data.Composite.Right;
                // here I have changed code order
                //code = (code << 1) | 1;
                //++bitsCount;
                code |=  1<<bitsCount++;
            }
            // if composite node, save position in the stack and go left
            else {
                bitsCount++;
                parents.push(current);
                current = CodesTree[current].Data.Composite.Left;
                //code <<= 1;
            }
        } // end while

		/*
		 * Setting decoding start jumps.
         */
        StartJumpTable.clear();
        const ui32 maxCode = 1<<DECODE_JUMP_LEN;
        StartJumpTable.resize(maxCode);
        // for each code with len equal to max jump len
//        for (ui32 code = 0; code < maxCode; code++) {
//            ui32 currentNode = TreeHead;
//            ui32 bit;
//            // seek for a leaf
//			  for (bit = 0; bit < DECODE_JUMP_LEN && !IsNodeLeaf(currentNode); bit++) {
//                const bool  dir = code & 1<<bit;
//                currentNode = dir ? CodesTree[currentNode].Data.Composite.Right : CodesTree[currentNode].Data.Composite.Left;
//            }
//            // set code remains len
//            // if
//			StartJumpTable[code].Len = DECODE_JUMP_LEN - bit;
//            // set ending node
//            StartJumpTable[code].Node = currentNode;
//        }
        // for each code with len equal to max jump len
        for (ui32 code = 0; code < maxCode; code++) {
            TNodeID currentNode = TreeHead;
            ui32 usedBits = 0;
            ui32 nodesSaved = 0;
            StartJumpTable[code].Node = TreeHead;
            StartJumpTable[code].Node2 = TreeHead;

            // seek for a leaf
            for (ui32 bit = 0; bit < DECODE_JUMP_LEN ; bit++) {
                const bool dir = code & 1<<bit;
                // move through the tree
                currentNode = dir 	? 	CodesTree[currentNode].Data.Composite.Right
                                    : 	CodesTree[currentNode].Data.Composite.Left;
                // if we found leaf
                if (IsNodeLeaf(currentNode) ) {
                    // if we could save the leaf node
                    if (nodesSaved < TStartJumpResult::SIZE) {
                        // set leaf node
                        switch (nodesSaved) {
                            case 0:
                                StartJumpTable[code].Node = currentNode;
                                break;
                            case 1:
                                StartJumpTable[code].Node2 = currentNode;
                                break;
                        }
                        // increment len of node queue
                        nodesSaved++;
                        // fix last used bits count
                        usedBits = bit + 1;
                        // reset position on the tree to the head
                        currentNode = TreeHead;
                    } else break;
                }
            }
            // get count of unused bits
            StartJumpTable[code].Len = DECODE_JUMP_LEN - usedBits;
            // if we reached the end of the code but have not saved any leaf
            if (StartJumpTable[code].Len == DECODE_JUMP_LEN) {
            // set ending node
                StartJumpTable[code].Node = currentNode;
                // increment len of node queue
                nodesSaved++;
                // reset unused bits
                StartJumpTable[code].Len = 0;
            }
        }
    }
    void PrintNode(TNodeID node) {
        char symb_str[sizeof(T)*8+1];
        char code_str[1<<sizeof(T)*8];
        cout <<"node: "<<node;
        if (!IsNodeLeaf(node))
            cout<< " right: " << CodesTree[node].Data.Composite.Right << " left: "<< CodesTree[node].Data.Composite.Left<<endl;
        else {
            const ui64 symb =  CodesTree[node].Data.Code;
            const ui64 code = DirectIndex[symb].first;
            const ui32 len = DirectIndex[symb].second;
            CodeToStr(code, code_str, len);
            CodeToStr(symb, symb_str, sizeof(T)*8);
            cout<<" : symb : "<< symb << " : "  << symb_str << " : code : "<< code_str <<" : "<< len << endl;
        }
    }
    void PrintTree() {
        for (TNodeID node = 0; node < CodesTree.size(); node++) {
            PrintNode(node);
        }
    }
    bool IsNodeLeaf(TNodeID node) {
        return node < TreeHead;
    }
};

/* Table of frequencies generator with two modes.
 * 1. A generation from raw file by simple counting.
 * 2. A generation from previously setted and written table.
 */
template <typename T>
class TFreqTable {	
	typedef ui32 TNumOfElements;

    TFreqsMap freqs;
public:
    enum TCreationModeRaw { CREATION_MODE_RAW = 0 };
    enum TCreationModeSet { CREATION_MODE_SET = 1 };
    /* This constructor read file from current position expecting there
     * saved frequency table. It will not rewind file stream afterwards.
     */
    template <typename TSource>
    TFreqTable(TTDataReader<TSource>& file, const TCreationModeSet) {
        SASSERT(sizeof(T) <= MAX_BYTES_IN_SYMBOL)

		TNumOfElements numElements=0;
		ui64 to_read = sizeof(TNumOfElements);
		ui64 readed = file.ReadInto(numElements);
        if (readed != to_read)
            throw TException(0, "TFreqTable::TFreqTable : Bad freq table file.");

        for (ui64 i=0; i< numElements; i++) {
			typename TFreqsMap::key_type first(0);
			typename TFreqsMap::mapped_type second(0);
			to_read = sizeof(typename TFreqsMap::key_type);
			readed = file.ReadInto(first);
            if (readed != to_read)
                throw TException(0, "TFreqTable::TFreqTable : Bad freq table file.");
			to_read = sizeof(typename TFreqsMap::mapped_type);
			readed = file.ReadInto(second);
            if (readed != to_read)
                throw TException(0, "TFreqTable::TFreqTable : Bad freq table file.");
            freqs[first] = second;
        }
DEBUG_CODE (Print();)

    }
    /* This constructor read file from current position till the end and counts
     * frequencies of each symbol. Finally it sets stream position at the begining.
     */
    template <typename TSource>
    TFreqTable(TTDataReader<TSource>& file, const TCreationModeRaw) {
        SASSERT(sizeof(T) <= MAX_BYTES_IN_SYMBOL)

		union {
            TFreqsMap::key_type asKey;
            T asT;
		} symb;
		symb.asKey = 0;

        while (true) {
            symb.asKey = 0;
			const ui64 readed = file.ReadInto(symb.asT);
			if (readed < sizeof(T))
                break;
            freqs[symb.asKey]++;
        }
        // todo check it if it works
        freqs[THuffmanCodec<T>::ESCAPE_SYMBOL] = 1;
        freqs[THuffmanCodec<T>::INVALID_CODE] = 1;

#ifdef DEBUG
        cout << "TFreqTable(ifstream *file, const TCreationModeRaw)" << endl;
        Print();
#endif
    }
    void Print() {
        const int numElements = freqs.size();
        printf("numElements: %d\n", numElements);
		for (typename TFreqsMap::iterator it=freqs.begin(); it!=freqs.end(); it++) {
			cout<<it->first<<","<< it->second<<endl;
        }
    }
    const TFreqsMap& GetFreqs() {
        return freqs;
    }
    template <typename TStorage>
    void Write(TTDataWriter<1, TStorage>& file) {
		const TNumOfElements numElements = freqs.size();
		file.WriteWhole(numElements);
		for (typename TFreqsMap::iterator it=freqs.begin(); it!=freqs.end(); it++) {
			file.WriteWhole(it->first);
			file.WriteWhole(it->second);
        }
    }
};

template <typename T, typename TIn, typename TOut>
void HuffmanFileCompress(TIn in, TOut out) {
        TTDataReader<TIn> inputDataReader(in);
        //TTBitReader< TTDataReader<TIn> > inputBitReader(in);
        TTBitWriter< 1, TTDataWriter<1, TOut> > output(out);
        TFreqTable<T> freqs(inputDataReader, TFreqTable<T>::CREATION_MODE_RAW);
        freqs.Write(output);
        inputDataReader.Rewind();
        THuffmanCodec<T> huffmanCodec(freqs.GetFreqs());
        huffmanCodec.Code(inputDataReader, output);
}

template <typename T, typename TIn, typename TOut>
void HuffmanFileDecompress(TIn in, TOut out) {
        //TTDataReader<TIn> inputDataReader(in);
        TTBitReader< TTDataReader<TIn> > inputBitReader(in);
        TTDataWriter<1, TOut> output(out);
        //TTBitWriter< 1, TTDataWriter<1, TOut> > output(out);
        TFreqTable<T> freqs(inputBitReader, TFreqTable<T>::CREATION_MODE_SET);
		THuffmanCodec<T> huffmanCodec(freqs.GetFreqs());
        huffmanCodec.Decode(inputBitReader, output);
}
/*r
template <typename T>
void HuffmanTestDecoding(FILE* in, FILE *out) {
        TTBitReader< TTDataReader<FILE*> > input(in);
        TTBitWriter< 1, TTDataWriter<1, FILE*> > output(out);
        TFreqTable<T> freqs(input, TFreqTable<T>::CREATION_MODE_SET);
        THuffmanCodec<T> huffmanCodec(freqs.GetFreqs());
		vector<char> decoded;
        const ui32 size = input.Size()*2;
		decoded.reserve(size);
		TTBitWriter< 1, TTDataWriter<1, vector<char>*> > vectorOut(&decoded);
//TIMER_CODE (TStopwatch watch("DECODING");)
        huffmanCodec.Decode(input, vectorOut);
//TIMER_CODE (watch.Print();)
}

template <typename T>
void HuffmanOldDecoding(FILE* fileIn, FILE *fileOut) {
		TIMER_CODE (TStopwatch watch("DECODING");)

		TTDataReader<FILE*> dataReader(fileIn);
		TFreqTable<T> freqs(dataReader, TFreqTable<T>::CREATION_MODE_SET);
		THuffmanCodec<T> huffmanCodec(freqs.GetFreqs());

		const ui64 MAX_BUF_SIZE  = 4096;
		vector<char> decoded;

		while (true) {
			const ui64 fileInSize = dataReader.Size();
			const bool canAtOnce = fileInSize <= MAX_BUF_SIZE;
			const ui64 toRead = canAtOnce ? fileInSize : MAX_BUF_SIZE;
			char buf[MAX_BUF_SIZE];
			const int readed = fread(buf, 1, toRead, fileIn);
			if (readed <= 0)
				break;

			const char* begin = buf;
			const char* end = buf+(readed-1);
//            TOldSchoolBitInput<1> oldSchoolBitInput(&begin, &end);
			decoded.reserve(fileInSize*2);

			try {
				huffmanCodec.oldDecode(buf, buf+readed-1, &decoded);
			} catch (...) {
				cout<<"Old Decoding fails"<<endl;
			}
		}

		TTDataWriter<1, FILE*> dataWriter(fileOut);
		dataWriter.WriteData((const char*)&decoded[0], decoded.size()-1);

		TIMER_CODE (watch.Print();)
}
*/
